import torch

from util.seed import set_seed
from defense.get_defense import get_defense

def main(model_name, defense_name, defense_config, suffix = ''):
    # 设置随机数种子
    set_seed()

    # 加载防御
    # defense_config = defense_config.copy()
    # defense_config['suffix'] = suffix
    defense = get_defense(defense_name, model_name, defense_config)

    # 对于需要重新训练的防御，进行训练，保存防御后的网络参数
    if defense.retrain:
        state_dict = defense.defend()
        path = 'data/enhanced_models/' + model_name + '_' + defense_name + suffix + '.pt'
        torch.save(state_dict, path)
        print('Model saved to ' + path)

    # 进行测试
    if defense.retrain:
        # 对于需要重新训练的防御，加载防御后的网络参数
        defense.model.load_state_dict(state_dict)
    else:
        # 对于不需要重新训练的防御，加载最初训练的网络参数
        defense.model.load_raw_state()
    defense.test(defense.model.get_test_loader(batch_size = 200, shuffle = False))
